#include "ARAPBuilder.h"

namespace ARAP
{
	const float MIN_DEFAULT = FLT_MAX;
	const float MAX_DEFAULT = -FLT_MAX;

	ARAPBuilder::ARAPBuilder(Mesh* mesh, const std::vector<LSCM::Chart>& charts, double lambda)
		: m_mesh(mesh), m_charts(charts), m_lambda(lambda)
	{
		m_mesh->add_property(m_Lts);
		m_mesh->add_property(m_halfedgeData);
		m_mesh->add_property(m_ids);
		m_mesh->add_property(m_chartIds);
	}

	ARAPBuilder::~ARAPBuilder()
	{
		m_mesh->remove_property(m_Lts);
		m_mesh->remove_property(m_halfedgeData);
		m_mesh->remove_property(m_ids);
		m_mesh->remove_property(m_chartIds);
	}

	void ARAPBuilder::Build(int times)
	{
		// 标记图表的面
		for (unsigned int i = 0; i < m_charts.size(); i++)
		{
			auto& chart = m_charts[i];
			for (auto& face : chart.faces())
			{
				// 面标记为所属图表的 Id
				m_mesh->property(m_chartIds, face) = chart.id();
			}
		}

		// 开始计算
		for (unsigned int i = 0; i < m_charts.size(); i++)
		{
			// 初始化数据
			auto& chart = m_charts[i];
			SparseMatrix MT = initializeData(chart);

			int t = 0;
			while (t < times)
			{
				localPhase(chart);
				globalPhase(chart, MT);
				t++;
			}

			// 保存并规范化
			storeUVs(chart);
		}
	}

	SparseMatrix ARAPBuilder::initializeData(const LSCM::Chart& chart)
	{
		// 计算每个面的信息并保存
		auto& faces = chart.faces();
		for (auto& face : faces)
		{
			int i = 0;
			auto h_it = m_mesh->fh_ccwbegin(face), h_end = m_mesh->fh_ccwend(face);

			// 记录句柄和方向
			Mesh::HalfedgeHandle h[3];
			Mesh::Point v[3];
			Mesh::Point dir[3];
			for (; h_it != h_end; h_it++, i++)
			{
				h[i] = *h_it;
				auto h1 = m_mesh->from_vertex_handle(h[i]);
				auto h2 = m_mesh->to_vertex_handle(h[i]);

				v[i] = m_mesh->point(h1);
				dir[i] = (m_mesh->point(h2) - v[i]).normalized();
			}

			// 计算余切
			double cot[3];
			for (int i = 0; i < 3; i++)
			{
				int j = (i + 1) % 3;
				double cosv = dir[i] | dir[j];
				cot[i] = 1 / std::tan(std::acos(-cosv));
			}

			// 计算参数坐标
			const auto v10 = v[1] - v[0];
			const auto v10Length = v10.length();
			const auto v20 = v[2] - v[0];

			const auto x = v10 / v10Length;
			const auto z = (x % v20).normalized();
			const auto y = z % x;

			v[0] = Mesh::Point(0, 0, 0);
			v[1] = Mesh::Point(v10Length, 0, 0);
			v[2] = Mesh::Point(v20 | x, v20 | y, 0);

			// 保存信息到半边
			for (int i = 0; i < 3; i++)
			{
				int j = (i + 1) % 3;

				auto dx = v[i] - v[j];
				auto& data = m_mesh->property(m_halfedgeData, h[i]);
				data.cot = cot[j];
				data.dx[0] = dx[0];
				data.dx[1] = dx[1];
			}
		}

		// 顶点编号
		auto& vertices = chart.vertices();
		int size = vertices.size();
		for (int i = 0; i < size; i++)
		{
			auto& v = vertices[i];
			m_mesh->property(m_ids, v) = i;
		}

		// 计算初始矩阵，并保存其预分解形式
		TripletList tripleList;
		for (int i = 0; i < size; i++)
		{
			auto& vi = vertices[i];
			int row = m_mesh->property(m_ids, vi);

			// 查找指向外侧的半边
			auto oh_it = m_mesh->voh_begin(vi), oh_end = m_mesh->voh_end(vi);
			for (; oh_it != oh_end; oh_it++)
			{
				// 跳过边界半边
				if (oh_it->is_boundary())
					continue;

				// 排除不属于图表的半边
				auto& face = m_mesh->face_handle(*oh_it);
				if (m_mesh->property(m_chartIds, face) != chart.id())
					continue;

				auto& vj = m_mesh->to_vertex_handle(*oh_it);
				int col = m_mesh->property(m_ids, vj);
				auto& data = m_mesh->property(m_halfedgeData, *oh_it);

				tripleList.push_back({ row, row, data.cot });
				tripleList.push_back({ row, col, -data.cot });
			}

			// 查找指向内侧的半边
			auto ih_it = m_mesh->vih_begin(vi), ih_end = m_mesh->vih_end(vi);
			for (; ih_it != ih_end; ih_it++)
			{
				// 跳过边界半边
				if (ih_it->is_boundary())
					continue;

				// 排除不属于图表的半边
				auto& face = m_mesh->face_handle(*ih_it);
				if (m_mesh->property(m_chartIds, face) != chart.id())
					continue;

				auto& vj = m_mesh->from_vertex_handle(*ih_it);
				int col = m_mesh->property(m_ids, vj);
				auto& data = m_mesh->property(m_halfedgeData, *ih_it);

				tripleList.push_back({ row, row, data.cot });
				tripleList.push_back({ row, col, -data.cot });
			}
		}

		std::cout << "Chart vertices num: " << size << std::endl;

		SparseMatrix M(size, size);
		M.setFromTriplets(tripleList.begin(), tripleList.end());

		// 构造加权系数矩阵
		TripletList diag;
		for (int i = 0; i < size; i++)
			diag.push_back({ i, i, m_lambda * size });
		SparseMatrix D(size, size);
		D.setFromTriplets(diag.begin(), diag.end());

		SparseMatrix MT = M.transpose();
		SparseMatrix MTM = MT * M + D;
		MTM.makeCompressed();

		m_solver.compute(MTM);
		auto info = m_solver.info();
		switch (info)
		{
		case Eigen::Success:
			std::cout << "Factorize successfully." << std::endl;
			break;
		case Eigen::NumericalIssue:
			std::cout << "Factorize NumericalIssue." << std::endl;
			break;
		case Eigen::NoConvergence:
			std::cout << "Factorize NoConvergence." << std::endl;
			break;
		case Eigen::InvalidInput:
			std::cout << "Factorize InvalidInput." << std::endl;
			break;
		default:
			break;
		}

		return MT;
	}

	void ARAPBuilder::localPhase(const LSCM::Chart& chart)
	{
		auto& faces = chart.faces();
		for (auto& face : faces)
		{
			MatrixX S(2, 2);
			S.setZero();

			auto h_it = m_mesh->fh_ccwbegin(face), h_end = m_mesh->fh_ccwend(face);
			for (; h_it != h_end; h_it++)
			{
				auto& data = m_mesh->property(m_halfedgeData, *h_it);
				auto& v1 = m_mesh->from_vertex_handle(*h_it);
				auto& v2 = m_mesh->to_vertex_handle(*h_it);

				MatrixXx1 u(2);
				auto uv = m_mesh->texcoord2D(v1) - m_mesh->texcoord2D(v2);
				u[0] = uv[0];
				u[1] = uv[1];

				S = S + u * data.dx.transpose() * data.cot;
			}

			// 奇异值分解，然后更新 Lt
			Eigen::JacobiSVD<MatrixX> svd(S, Eigen::ComputeFullU | Eigen::ComputeFullV);
			m_mesh->property(m_Lts, face) = svd.matrixU() * svd.matrixV().transpose();
		}
	}

	void ARAPBuilder::globalPhase(const LSCM::Chart& chart, const SparseMatrix& M)
	{
		auto& vertices = chart.vertices();
		int size = vertices.size();
		MatrixX b(size, 2);
		b.setZero();

		for (int i = 0; i < size; i++)
		{
			auto& vi = vertices[i];
			int row = m_mesh->property(m_ids, vi);

			// 查找指向外侧的半边
			auto oh_it = m_mesh->voh_begin(vi), oh_end = m_mesh->voh_end(vi);
			for (; oh_it != oh_end; oh_it++)
			{
				// 排除掉边界半边
				if (oh_it->is_boundary())
					continue;

				// 排除不属于图表的半边
				auto& face = m_mesh->face_handle(*oh_it);
				if (m_mesh->property(m_chartIds, face) != chart.id())
					continue;

				auto& vj = m_mesh->to_vertex_handle(*oh_it);
				int col = m_mesh->property(m_ids, vj);
				auto& data = m_mesh->property(m_halfedgeData, *oh_it);

				MatrixX& Lt = m_mesh->property(m_Lts, face);
				MatrixXx1 db = Lt * data.dx * data.cot;

				// 设置系数
				b(row, 0) += db[0];
				b(row, 1) += db[1];
			}

			// 查找指向内侧的半边
			auto ih_it = m_mesh->vih_begin(vi), ih_end = m_mesh->vih_end(vi);
			for (; ih_it != ih_end; ih_it++)
			{
				// 排除掉边界半边
				if (ih_it->is_boundary())
					continue;

				// 排除不属于图表的半边
				auto& face = m_mesh->face_handle(*ih_it);
				if (m_mesh->property(m_chartIds, face) != chart.id())
					continue;

				auto& vj = m_mesh->from_vertex_handle(*ih_it);
				int col = m_mesh->property(m_ids, vj);
				auto& data = m_mesh->property(m_halfedgeData, *ih_it);

				MatrixX& Lt = m_mesh->property(m_Lts, face);
				MatrixXx1 db = -Lt * data.dx * data.cot;

				// 设置系数
				b(row, 0) += db[0];
				b(row, 1) += db[1];
			}
		}

		// 求解全局系统
		MatrixX x = m_solver.solve(M * b);
		for (int i = 0; i < size; i++)
		{
			auto& vi = vertices[i];
			Mesh::TexCoord2D tex;
			tex[0] = x(i, 0);
			tex[1] = x(i, 1);

			// 更新参数
			m_mesh->set_texcoord2D(vi, tex);
		}
	}

	void ARAPBuilder::storeUVs(const LSCM::Chart& chart)
	{
		// 初始化最大最小纹理坐标
		auto minUV = Mesh::TexCoord2D(MIN_DEFAULT, MIN_DEFAULT);
		auto maxUV = Mesh::TexCoord2D(MAX_DEFAULT, MAX_DEFAULT);

		for (const auto& vertex : chart.vertices())
		{
			// 固定顶点已经有参数坐标，只需要考虑活动顶点
			auto uv = m_mesh->texcoord2D(vertex);

			// 记录最大、最小纹理坐标
			minUV.minimize(uv);
			maxUV.maximize(uv);
		}

		auto diff = maxUV - minUV;
		const auto length = diff.max();

		// 规范化纹理坐标
		for (const auto& vertex : chart.vertices())
		{
			auto uv = m_mesh->texcoord2D(vertex);
			const auto nuv = (uv - minUV) / length;

			// 记录顶点的纹理坐标
			m_mesh->set_texcoord2D(vertex, nuv);
		}
	}

} // namespace ARAP
